//
//  Integrator.cpp
//  RayTracer
//
//  Created by 俞云康 on 2/15/16.
//  Copyright © 2016 yushroom. All rights reserved.
//

#include "Integrator.hpp"
#include "Renderer.hpp"
#include "Sample.hpp"
#include "Params.hpp"
#include "Material.hpp"
#include "Primitive.hpp"
#include <vector>
using namespace std;

shared_ptr<SurfaceIntegrator> SurfaceIntegrator::create(const Params& params) {
    shared_ptr<SurfaceIntegrator> si;
    string type = params.getType();
    if (type == "direct" || type == "directLight") {
		int maxDepth = params.getIntAttribute("maxDepth", 5);
        si = make_shared<DirectLightingIntegrator>(maxDepth);
	} else if (type == "path") {
		si = make_shared<PathIntegrator>();
	} else if (type == "normal") {
        si = make_shared<NormalIntegrator>();
    } else if (type == "depth") {
        si = make_shared<DepthIntegrator>();
    } else if (type == "ambientOcclusion") {
        si = make_shared<AmbientOcclusion>();
	} else if (type == "whitted") {
		si = make_shared<WhittedIntegrator>();
	} else {
        error("unknown SurfaceIntegrator type: %s\n", type.c_str());
    }
    return si;
}


Color SpecularReflect(const Ray& ray, const BSDF& bsdf, RNG& rng, const HitInfo& hit, const Renderer& renderer, const Scene& scene, Sampler& sampler) {
    Direction wo = -ray.dir, wi;
    float pdf = 0.f;
    auto& p = hit.position;
    auto& n = hit.normal;
    Color f = bsdf.sample_f(wo, &wi, random_sample(rng), &pdf, BxDFType(BxDF_REFLECTION | BxDF_SPECULAR));
    //auto pp = Normalize(wi + wo) - n;
    //Assert(pp == float3(0, 0, 0));
    Color L(0, 0, 0);
    if (pdf > 0.f && !is_black(f) && AbsDot(wi, n) != 0.f) {
        Ray r(p, wi);
        r.depth = ray.depth + 1;
        Color Li =  renderer.Li(scene, r, sampler, rng);
        L = f*Li*AbsDot(wi, n) / pdf;
    }
    return L;
}

Color SpecularTransmit(const Ray& ray, const BSDF& bsdf, RNG& rng, const HitInfo& hit, const Renderer& renderer, const Scene& scene, Sampler& sampler) {
	Direction wo = -ray.dir, wi;
	float pdf = 0.f;
	auto& p = hit.position;
	auto& n = hit.normal;
	Color f = bsdf.sample_f(wo, &wi, random_sample(rng), &pdf, BxDFType(BxDF_TRANSMISSION | BxDF_SPECULAR));
	//auto pp = Normalize(wi + wo) - n;
	//Assert(pp == float3(0, 0, 0));
	Color L(0, 0, 0);
	if (pdf > 0.f && !is_black(f) && AbsDot(wi, n) != 0.f) {
		Ray r(p, wi);
		r.depth = ray.depth + 1;
		Color Li = renderer.Li(scene, r, sampler, rng);
		L = f*Li*AbsDot(wi, n) / pdf;
	}
	return L;
}

Color NormalIntegrator::Li(const Scene& scene, const Renderer& renderer, const Ray& ray, const HitInfo& hit, Sampler& sampler, RNG& rng) const
{
    return Color(hit.normal.x, hit.normal.y, hit.normal.z) * 0.5f + 0.5f;
}

Color DepthIntegrator::Li(const Scene& scene, const Renderer& renderer, const Ray& ray, const HitInfo& hit, Sampler& sampler, RNG& rng) const
{
    float max_depth = 0;
    for (int i = 0; i < 8; ++i) {
        Position corner = scene.worldBounds().getCorner(i);
        float dist = Distance(ray.o, corner);
        if (max_depth < dist)
            max_depth = dist;
    }
    float c = 1.0f - ray.maxt / max_depth;
    return Color(c, c, c);
}

Color AmbientOcclusion::Li(const Scene& scene, const Renderer& renderer, const Ray& ray, const HitInfo& hit, Sampler& sampler, RNG& rng) const
{
    const float3 p = hit.position + EPS * hit.normal;
    float n_clear = 0;
    uint32_t scramble[2] = { rng.UniformUInt32() , rng.UniformUInt32() };
    float u[2];
    for (int i = 0; i < n_samples; ++i) {
        Sample02(i, scramble, u);
        float3 dir = UniformSampleSphere(u[0], u[1]);
        if (Dot(dir, hit.normal) < 0) dir = -dir;
        Ray r(p, dir);
        r.maxt = max_t;
        if (!scene.intersect(r))
            n_clear++;
    }
    float c = float(n_clear) / float(n_samples);
    return Color(c, c, c);
}


inline float PowerHeuristic(int nf, float fpdf, int ng, float gpdf) {
    float f = nf * fpdf, g = ng * gpdf;
    f *= f;
    g *= g;
    return f / (f + g);
}

Color EstimateDirect(const Scene& scene, const Renderer& renderer, const AreaLight& light, const Position& p, const Normal& n, const Direction& wo, const BSDF& bsdf, const Sample& sample, RNG& rng, BxDFType flags) {
    Color Ld(0, 0, 0);
    // Sample light source with multiple importance sampling
    Direction wi;
    float light_pdf, bsdf_pdf;
    VisibilityTester vis;
    Color Li = light.sample_L(p, sample, &wi, &light_pdf, &vis);
    if (light_pdf > 0 && !is_black(Li)) {
        // Add light's contribution to reflected radiance
        Color f = bsdf.f(wo, wi, flags);
        if (!is_black(f) && vis.unocclude(scene)) {
            bsdf_pdf = bsdf.pdf(wo, wi, flags);
            float weight = PowerHeuristic(1, light_pdf, 1, bsdf_pdf);
            Ld += f * Li * (AbsDot(wi, n) * weight / light_pdf);
        }
    }
    
    // Sample BSDF with multiple importance sampling
    BxDFType sampled_type;
    Color f = bsdf.sample_f(wo, &wi, sample, &bsdf_pdf, flags, &sampled_type);
    if (!is_black(f) && bsdf_pdf > 0) {
        float weight = 1.f;
        if (!(sampled_type & BxDF_SPECULAR)) {
            light_pdf = light.pdf(p, wi);
            if (light_pdf == 0)
                return Ld;
            weight = PowerHeuristic(1, bsdf_pdf, 1, light_pdf);
        }
        
        // Add light contribution from BSDF sampling
        HitInfo hit;
        Color Li(0, 0, 0);
        Ray ray(p, wi);
        if (scene.intersect(ray, &hit)) {
            if (hit.primitive->getAreaLight() == &light) {
                Li = light.L(hit.position, hit.normal, -wi);
            }
        }
        if (!is_black(Li)) {
            Ld += f * Li * AbsDot(wi, n) * weight / bsdf_pdf;
        }
    }
    return Ld;
}


Color UniformSampleAllLights(const Scene& scene, const Renderer& renderer, const Position& p, const Normal& n, const Direction& wo, const BSDF& brdf, Sampler& sampler, size_t num_samples, RNG& rng) {
    Color L(0, 0, 0);
    for (auto& light : scene.area_lights) {
        Color Ld(0, 0, 0);
        for (size_t j = 0; j < num_samples; ++j) {
            Sample sample = sampler.get2D();
            Ld += EstimateDirect(scene, renderer, *light, p, n, wo, brdf, sample, rng, BxDFType(BxDF_ALL & ~BxDF_SPECULAR));
        }
        L += Ld / (const float)num_samples;
    }
    return L;
}

Color UniformSampleOneLight(const Scene& scene, const Renderer& renderer, const Position& p, const Normal& n, const Direction& wo, const BSDF& brdf, Sampler& sampler, size_t num_samples, RNG& rng) {
	Color L(0, 0, 0);
	int nLights = int(scene.area_lights.size());
	if (nLights == 0)
		return L;
	int lightNum = floor2Int(rng.UniformFloat() * nLights);
	lightNum = min(lightNum, nLights - 1);
	auto light = scene.area_lights[lightNum];

	Sample sample = sampler.get2D();
	return (float)nLights * EstimateDirect(scene, renderer, *light, p, n, wo, brdf, sample, rng, BxDFType(BxDF_ALL & ~BxDF_SPECULAR));
}

Color PathIntegrator::Li(const Scene& scene, const Renderer& renderer, const Ray& ray, const HitInfo& hit, Sampler& sampler, RNG& rng) const
{
	Color pathThroughput(1.f, 1.f, 1.f);
	Color L(0.f, 0.f, 0.f);
	Ray r(ray);
	bool specularBounce = false;
	HitInfo localIsect;
	const HitInfo *isectp = &hit;
	for (int bounces = 0; ; ++bounces) {
		if (bounces == 0 || specularBounce)
			L += pathThroughput * isectp->Le(-r.dir);

		auto bsdf = isectp->material->getBSDF(*isectp);
		const Position &p = isectp->position;
		const Normal &n = isectp->normal;

		Direction wo = -r.dir;
		L += pathThroughput * 
			UniformSampleOneLight(scene, renderer, p, n, wo, *bsdf, sampler, sampler.m_samplersPerPixel, rng);

		Sample outgoingBSDFSample{ rng.UniformFloat(), rng.UniformFloat() };
		Direction wi;
		float pdf;
		BxDFType flags;
		Color f = bsdf->sample_f(wo, &wi, outgoingBSDFSample, &pdf, BxDF_ALL, &flags);
		if (is_black(f) || pdf == 0.f) {
			break;
		}
		specularBounce = (flags & BxDF_SPECULAR) != 0;
		pathThroughput *= f * AbsDot(wi, n) / pdf;
		r = Ray(p, wi);

		if (bounces > 3) {
			float continueProbability = std::min(.5f, pathThroughput.y);
			if (rng.UniformFloat() > continueProbability)
				break;
			pathThroughput /= continueProbability;
		}
		if (bounces == m_maxDepth)
			break;

		if (!scene.intersect(r, &localIsect)) {
			//if (specularBounce) {
			//	//for (size_t i = 0; i < scene.area_lights.size(); ++i)
			//	for (auto& al : scene.area_lights)
			//		L += pathThroughput * al->Le(r);
			//}
			break;
		}
		//pathThroughput *= renderer->Transmittance(scene, ray, NULL, rng, arena);
		isectp = &localIsect;
	}
	return L;
}

Color DirectLightingIntegrator::Li(const Scene& scene, const Renderer& renderer, const Ray& ray, const HitInfo& hit, Sampler& sampler, RNG& rng) const
{
    Color L(0, 0, 0);
    auto& n = hit.normal;
    auto p = hit.position + EPS * n;
    auto wo = -ray.dir;
    if (hit.primitive->getAreaLight() != nullptr) {
        return hit.primitive->getAreaLight()->L(p, n, wo);
    }
	//L += hit.Le(wo);

    //auto brdf = hit.material->get_brdf(hit);
    auto bsdf = hit.material->getBSDF(hit);
    L += UniformSampleAllLights(scene, renderer, p, n, wo, *bsdf, sampler, sampler.m_samplersPerPixel, rng);
    if (ray.depth + 1 < m_maxDepth) {
        L += SpecularReflect(ray, *bsdf, rng, hit, renderer, scene, sampler);
		L += SpecularTransmit(ray, *bsdf, rng, hit, renderer, scene, sampler);
    }
    return L;
}

Color DirectIlluminationIngetrator::Li(const Scene& scene, const Renderer& renderer, const Ray& ray, const HitInfo& hit, Sampler& sampler, RNG& rng) const
{
    Color L(0, 0, 0);
    auto& n = hit.normal;
    auto p = hit.position + EPS * n;
    auto wo = -ray.dir;
    //if (hit.primitive->getAreaLight() != nullptr) {
    //    return hit.primitive->getAreaLight()->L(p, n, wo);
    //}
	L += hit.Le(wo);

    //auto brdf = hit.material->get_brdf(hit);
    auto bsdf = hit.material->getBSDF(hit);
    L += UniformSampleAllLights(scene, renderer, p, n, wo, *bsdf, sampler, 4, rng);
    return L;
}

Color WhittedIntegrator::Li(const Scene& scene, const Renderer& renderer, const Ray& ray, const HitInfo& hit, Sampler& sampler, RNG& rng) const
{
	Color L(0, 0, 0);
	auto bsdf = hit.material->getBSDF(hit);
	const Position &p = hit.position;
	const Normal &n = hit.normal;
	Direction wo = -ray.dir;

	L += hit.Le(wo);

	for (auto light : scene.area_lights) {
		Direction wi;
		float pdf;
		VisibilityTester visibility;
		Color Li = light->sample_L(p, sampler.get2D(), &wi, &pdf, &visibility);
		if (is_black(Li) || pdf == 0.f) continue;
		Color f = bsdf->f(wo, wi);
		if (!is_black(f) && visibility.unocclude(scene))
			L += f * Li * AbsDot(wi, n);
	}

	if (ray.depth + 1 < m_maxDepth) {
		L += SpecularReflect(ray, *bsdf, rng, hit, renderer, scene, sampler);
		L += SpecularTransmit(ray, *bsdf, rng, hit, renderer, scene, sampler);
	}

	return L;
}